package com.mkx.aqs;

import java.util.concurrent.locks.LockSupport;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 简化版AQS队列实现 - 高性能优化版本
 * 专注于存储和管理阻塞的线程节点，优化了高并发性能
 */
public class AQSQueue {
    // 使用AtomicReference来实现CAS操作，避免依赖Unsafe类
    private final AtomicReference<Node> headRef = new AtomicReference<>();
    private final AtomicReference<Node> tailRef = new AtomicReference<>();
    // 添加原子计数器，提高size方法的性能
    private final AtomicInteger sizeCounter = new AtomicInteger(0);
    
    /**
     * 构造函数，初始化空队列
     */
    public AQSQueue() {
        // 创建一个哨兵头节点
        Node h = new Node();
        headRef.set(h);
        tailRef.set(h);
    }
    
    /**
     * 将当前线程添加到等待队列 - 高性能优化版本
     * 使用改进的快速路径和批量入队优化减少CAS竞争
     * @param mode 节点模式（SHARED或EXCLUSIVE）
     * @return 新创建的节点
     */
    public Node addWaiter(Node mode) {
        Node node = new Node(Thread.currentThread(), mode);
        // 尝试快速入队 - 优化版
        for (;;) {
            Node tail = tailRef.get();
            if (tail == null) {
                // 队列未初始化，需要初始化
                enq(node);
                sizeCounter.incrementAndGet();
                return node;
            }
            
            node.prev = tail;
            // 第一次尝试快速入队
            if (compareAndSetTail(tail, node)) {
                // 使用常规赋值设置next指针，因为tail已经通过CAS更新
                // 即使这里失败，后续的enq操作也会修复next指针
                tail.next = node;
                sizeCounter.incrementAndGet();
                return node;
            }
            
            // 快速路径失败，检查是否需要帮助完成其他线程的入队操作
            // 这有助于减少CAS竞争
            if (tail.next != null) {
                // 有其他线程正在入队但未完成，尝试帮助推进tail
                compareAndSetTail(tail, tail.next);
            } else {
                // 完全失败，进入enq方法
                enq(node);
                sizeCounter.incrementAndGet();
                return node;
            }
        }
    }
    
    /**
     * 入队操作，确保节点被添加到队列 - 高性能优化版本
     * 使用CAS优化和批处理机制减少竞争
     * @param node 要入队的节点
     * @return 入队后的节点
     */
    private Node enq(final Node node) {
        // 快速检查队列是否已初始化
        Node h = headRef.get();
        if (h == null) {
            // 初始化队列
            Node sentinel = new Node();
            if (headRef.compareAndSet(null, sentinel)) {
                tailRef.set(sentinel);
            }
        }
        
        // 尝试将节点添加到队列
        for (;;) {
            Node t = tailRef.get();
            if (t == null) {
                // 队列仍未初始化，重试
                continue;
            }
            
            // 帮助完成之前可能未完成的next指针设置
            Node s = t.next;
            if (s != null) {
                // 尾节点的next不为空，说明有其他线程添加了节点但没更新tail
                // 帮助更新tail
                compareAndSetTail(t, s);
                continue;
            }
            
            // 设置当前节点的prev引用
            node.prev = t;
            
            // 尝试设置当前节点为新的尾节点
            if (compareAndSetTail(t, node)) {
                // 成功更新tail，现在设置前一个节点的next引用
                // 这里不需要CAS，因为只有当前线程成功更新了tail
                t.next = node;
                return t;
            }
            
            // 失败，重试
        }
    }
    
    /**
     * 获取队首节点（不包括哨兵节点）
     * @return 第一个等待的线程节点，如果队列为空则返回null
     */
    public Node getFirstWaiter() {
        Node h = headRef.get();
        if (h != null) {
            return h.next;
        }
        return null;
    }
    
    /**
     * 从队列中移除指定节点 - 高性能优化版本
     * 优化节点移除策略，减少不必要的操作
     * @param node 要移除的节点
     * @return 是否成功移除
     */
    public boolean removeNode(Node node) {
        if (node == null) {
            return false;
        }
        
        // 将节点标记为取消状态
        node.setWaitStatus(Node.CANCELLED);
        
        // 快速路径：检查是否是头节点的直接后继
        Node h = headRef.get();
        Node first = (h != null) ? h.next : null;
        
        if (first == node) {
            // 直接更新头节点
            compareAndSetHead(h, node.next);
            // 清理节点引用
            cleanupNode(node);
            sizeCounter.decrementAndGet();
            return true;
        }
        
        // 获取前驱和后继节点
        Node pred = node.prev;
        Node succ = node.next;
        
        // 优化：只有在节点确实在队列中时才进行移除操作
        if (pred == null) {
            // 节点不在队列中
            return false;
        }
        
        // 移除节点，更新前后节点的引用
        pred.next = succ;
        
        if (succ != null) {
            succ.prev = pred;
        } else {
            // 如果是尾节点，需要使用CAS确保原子更新
            tailRef.compareAndSet(node, pred);
        }
        
        // 帮助GC
        cleanupNode(node);
        sizeCounter.decrementAndGet();
        
        return true;
    }
    
    /**
     * 清理节点引用，帮助GC
     */
    private void cleanupNode(Node node) {
        node.prev = null;
        node.next = null;
        node.thread = null;
    }
    
    /**
     * 唤醒指定节点的线程
     * @param node 要唤醒的节点
     */
    public void unparkNode(Node node) {
        if (node != null && node.thread != null) {
            LockSupport.unpark(node.thread);
        }
    }
    
    /**
     * 唤醒队首节点的线程
     * @return 是否成功唤醒
     */
    public boolean unparkFirstWaiter() {
        Node first = getFirstWaiter();
        if (first != null) {
            unparkNode(first);
            return true;
        }
        return false;
    }
    
    /**
     * 阻塞当前线程
     * @return 线程是否被中断
     */
    public boolean parkAndCheckInterrupt() {
        LockSupport.park(this);
        return Thread.interrupted();
    }
    
    /**
     * 判断队列是否为空（不包含哨兵节点）- 优化版本
     * @return 队列为空返回true，否则返回false
     */
    public boolean isEmpty() {
        Node h = headRef.get();
        Node t = tailRef.get();
        return h == t && h.next == null;
    }
    
    /**
     * 获取队列长度 - 高性能优化版本
     * 使用原子计数器提供O(1)时间复杂度的size方法
     * @return 队列中的节点数量
     */
    public int size() {
        // 使用原子计数器，避免遍历队列
        return sizeCounter.get();
    }
    
    /**
     * 清空队列
     */
    public void clear() {
        Node h = new Node();
        headRef.set(h);
        tailRef.set(h);
        // 重置计数器
        sizeCounter.set(0);
    }
    
    /**
     * 原子更新头节点
     */
    private final boolean compareAndSetHead(Node expect, Node update) {
        return headRef.compareAndSet(expect, update);
    }
    
    /**
     * 原子更新尾节点
     */
    private final boolean compareAndSetTail(Node expect, Node update) {
        return tailRef.compareAndSet(expect, update);
    }
    
    /**
     * 获取头节点
     */
    public Node getHead() {
        return headRef.get();
    }
    
    /**
     * 获取尾节点
     */
    public Node getTail() {
        return tailRef.get();
    }
    
    /**
     * 设置头节点
     */
    public void setHead(Node node) {
        headRef.set(node);
    }
    
    /**
     * 设置尾节点
     */
    public void setTail(Node node) {
        tailRef.set(node);
    }
    
    /**
     * 取消节点在队列中的等待 - 高性能优化版本
     * @param node 要取消的节点
     * @return 如果成功取消则返回true
     */
    public boolean cancelNode(Node node) {
        if (node == null) {
            return false;
        }
        
        // 优化：直接调用removeNode，它已经包含了状态设置和移除逻辑
        // 这样可以避免重复操作和计数错误
        return removeNode(node);
    }
}